#include <mc_controller/dwa.h>
#include <sensor_msgs/point_cloud2_iterator.h>

DWAController::DWAController(){
    dwaPredictedTrajPub_ = nh_.advertise<sensor_msgs::PointCloud2>("/dwa_predicted_trajs", 1, true);
    dwaBestTrajPub_ = nh_.advertise<nav_msgs::Path>("/dwa_best_traj", 1, true);

    buildVelocitySpace();
    predictTraj();
}


DWAController::~DWAController(){

}

// 构建速度空间
void DWAController::buildVelocitySpace(){
    int linear_samples_nums = ceil((baseConfig_.max_vel_x - baseConfig_.min_vel_x) 
        / dwaConfig_.linear_sampling_resolution) + 1;
    int angular_samples_nums = ceil((baseConfig_.max_vel_yaw - baseConfig_.min_vel_yaw) 
        / dwaConfig_.angular_sampling_resolution) + 1;
    
    if(linear_samples_nums <=0 || angular_samples_nums <=0){
        ROS_ERROR("please check base config params!");
        linear_samples_nums = 1;
        angular_samples_nums = 1;
    }

    geometry_msgs::Twist vel;
    for(int i=0; i<linear_samples_nums; ++i){
        vel.linear.x = baseConfig_.min_vel_x + i * dwaConfig_.linear_sampling_resolution;
        if(fabs(vel.linear.x) < 0.01) continue;      //不计算0速
        if(vel.linear.x > baseConfig_.max_vel_x) break;
        for(int j=0; j<angular_samples_nums; ++j){
            vel.angular.z = baseConfig_.min_vel_yaw + j * dwaConfig_.angular_sampling_resolution;
            if(vel.angular.z > baseConfig_.max_vel_yaw) break;
            velSpace_.emplace_back(vel);
            
        }
    }
}

// 计算预测轨迹
void DWAController::predictTraj(){
    for(auto& vel : velSpace_){
        Trajectory traj;
        traj.cmd_vel = vel;
        // 计算预测点集数量
        int predicted_points = ceil(dwaConfig_.predicted_time / dwaConfig_.time_sampling_resolution) + 1;
        double x = 0.0, y=0.0, yaw=0.0;
        for(int i=0; i<predicted_points; ++i){
            if(i * dwaConfig_.time_sampling_resolution > dwaConfig_.predicted_time) break;
            // 根据运动模型，计算轨迹  此处为差速运动模型
            x += vel.linear.x * cos(yaw) * dwaConfig_.time_sampling_resolution;
            y += vel.linear.x * sin(yaw) * dwaConfig_.time_sampling_resolution;
            yaw += vel.angular.z * dwaConfig_.time_sampling_resolution;       
            // 消息转换
            geometry_msgs::PoseStamped curr_pose;
            curr_pose.pose.position.x = x;
            curr_pose.pose.position.y = y;
            curr_pose.pose.orientation = msg_convert::toGeometryQuaternion(0.0, 0.0, yaw);
            traj.waypoints.emplace_back(curr_pose);    
        }
        predictedTraj_.emplace_back(traj);
    }
    publishPointCloud2(predictedTraj_);
}

// 发布点云数据
void DWAController::publishPointCloud2(vector<Trajectory> predicted_trajs){
    vector<geometry_msgs::Point> waypoints;
    for(auto traj : predicted_trajs){
        for(auto pose_stamp : traj.waypoints){
            auto point = pose_stamp.pose.position;
            point.z = traj.cmd_vel.linear.x;  //此处按速度vx分层显示
            // cout << msg_convert::toString(pose_stamp.pose) << endl;
            waypoints.emplace_back(point);
        }
    }
    sensor_msgs::PointCloud2 cloud_msg;
    cloud_msg.header.frame_id = baseFrameId_;
    cloud_msg.header.stamp = ros::Time::now();
    cloud_msg.height = 1;
    cloud_msg.width = waypoints.size();
    cloud_msg.is_dense = false;

    sensor_msgs::PointCloud2Modifier modifier(cloud_msg);
    modifier.setPointCloud2FieldsByString(1, "xyz");
    modifier.resize(waypoints.size());

    // 使用迭代器填充点云消息的数据
    sensor_msgs::PointCloud2Iterator<float> iter_x(cloud_msg, "x");
    sensor_msgs::PointCloud2Iterator<float> iter_y(cloud_msg, "y");
    sensor_msgs::PointCloud2Iterator<float> iter_z(cloud_msg, "z");
    for(const auto& point : waypoints) {
        *iter_x = point.x;
        *iter_y = point.y;
        *iter_z = point.z;

        ++iter_x;
        ++iter_y;
        ++iter_z;
    }
    dwaPredictedTrajPub_.publish(cloud_msg);
}

// 发布最佳轨迹
void DWAController::publishBestTraj(Trajectory best_traj){
    nav_msgs::Path path_msg;
    path_msg.header.frame_id = baseFrameId_;
    path_msg.header.stamp = ros::Time::now();
    for(auto pose_stamp : best_traj.waypoints){
        path_msg.poses.emplace_back(pose_stamp);
    }
    dwaBestTrajPub_.publish(path_msg);
}

// 计算航向角度评分  轨迹末端点角度和末端点与目标点连线的角度差
double DWAController::headingScore(geometry_msgs::PoseStamped traj_end_point, 
    geometry_msgs::PoseStamped goal){
    // 末端点与目标点连线的角度theta
    double goal_angle = math_utils::angleBetweenPoses(traj_end_point.pose, goal.pose);
    // 轨迹末端点角度 与  theta 的角度差
    double theta = math_utils::diffAngle(
        msg_convert::toYaw(traj_end_point.pose.orientation),
        goal_angle);
    return 180.0 - fabs(math_utils::toDeg(theta));        
}

// 速度评分
double DWAController::velocityScore(geometry_msgs::Twist cmd_vel){
    return baseConfig_.min_vel_x - cmd_vel.linear.x;     //速度越大越好
}

// 计算障碍物评分  计算轨迹末端点距离最近障碍物距离的评价函数 碰撞检测
bool DWAController::obstaclesScore(geometry_msgs::PoseStamped traj_end_point, 
    geometry_msgs::Twist cmd_vel, 
    vector<geometry_msgs::PoseStamped>& obstacles, 
    double inflation_radius, double& min_distance){
    // 计算轨迹末端点距离最近障碍物最小距离
    min_distance = std::numeric_limits<float>::infinity();
    for(auto& pose_stamp : obstacles){
        double dis = math_utils::distanceBetweenPoses(traj_end_point.pose, pose_stamp.pose);
        dis -= inflation_radius;    //减去障碍物膨胀半径
        if(dis < 0){
            return false;
        }
        // cout << "[obstaclesCollision] dis: " << dis << endl;
        min_distance = min(min_distance, dis);
    }
    // cout << "[obstaclesCollision] cmd_vel: " << msg_convert::toString(cmd_vel) << ", min_distance: " << min_distance << endl;

    //没有障碍物时，限制最大值， 此处设置为3 * inflation_radius
    min_distance = min(min_distance, 3 * inflation_radius);     
    return true;     
}


// 碰撞检测  简单的距离计算，后期优化
bool DWAController::obstaclesCollision(vector<geometry_msgs::PoseStamped> waypoints, 
    vector<geometry_msgs::PoseStamped>& obstacles, double inflation_radius){
    for(auto pose_stamp : waypoints){
        for(auto pose_obstacle : obstacles){
            float dis = math_utils::distanceBetweenPoses(pose_stamp.pose, pose_obstacle.pose);
            if(dis < inflation_radius){
                return true;
            }
        }
    }
    return false;
}

// 查找最佳轨迹
bool DWAController::findBestTrajectory(Trajectory& best_traj, geometry_msgs::PoseStamped local_goal){
    double inflation_radius = 0.2;   //膨胀半径
    double lidar_valid_range = 3.0;  //有效距离

    vector<geometry_msgs::PoseStamped> obstacles;  
    // if(pointCloudXYZPtr_){
    //     for(auto point : pointCloudXYZPtr_->points){
    //         // 计算lidar_valid_range米内障碍物数据
    //         if(hypot(point.x, point.y) > lidar_valid_range) continue;
    //         obstacles.emplace_back(msg_convert::ToPose2D(point));
    //     }
    //     cout << "scan_data size: " << pointCloudXYZPtr_->points.size() 
    //         << ", obstacles size: " << obstacles.size() << endl; 
    // }

    vector<Trajectory> valid_trajs;         //有效轨迹
    Trajectory sum_traj;    //用于归一化处理
    int best_i = -1;
    double best_traj_score = std::numeric_limits<double>::infinity();
    for(int i=0; i<predictedTraj_.size(); ++i){
        auto traj = predictedTraj_[i];
        if(traj.waypoints.empty()) continue;
        if(obstaclesCollision(traj.waypoints, obstacles, inflation_radius)){
            continue;
        }
        geometry_msgs::PoseStamped end_waypoint = traj.waypoints.back();
        // traj.obstacle_score = obstaclesCollision(end_waypoint, traj.cmd_vel, obstacles, inflation_radius);
        if(!obstaclesScore(end_waypoint, traj.cmd_vel, obstacles, inflation_radius, traj.obstacle_score)) 
            continue;
        // 计算刹车距离
        double stop_distance = maxStopDistance(traj.cmd_vel, 1.0);
        // cout << "[obstaclesCollision] stop_distance: " << stop_distance << ",  traj.obstacle_score: " <<  traj.obstacle_score << endl;
        if(traj.obstacle_score < stop_distance){
            continue;
        }
        // // cout << "traj.obstacle_score: " << traj.obstacle_score << endl;
        // if(isinf(traj.obstacle_score)) continue;
        traj.goal_score = math_utils::distanceBetweenPoses(end_waypoint.pose, local_goal.pose);
        traj.speed_score = velocityScore(traj.cmd_vel);
        traj.heading_score = headingScore(end_waypoint, local_goal);
        traj.traj_score = traj.speed_score + traj.goal_score;
        if(traj.traj_score < best_traj_score){
            best_traj_score = traj.traj_score;
            best_i = i;
        }

        valid_trajs.emplace_back(traj);
        // cout << "[debug | findBestTrajectory] --> goal_score: " << traj.goal_score
        //     // << ", obstacle_score: " << traj.obstacle_score
        //     << ", heading_score=" << traj.heading_score
        //     << ", speed_score=" << traj.speed_score
        //     << ", traj_score=" << traj.traj_score
        //     << endl;
        // 计算总分
        sum_traj.obstacle_score += traj.obstacle_score;
        sum_traj.speed_score += traj.speed_score;
        sum_traj.heading_score += traj.heading_score;
    }

    // for(int i=0; i<valid_trajs.size(); ++i){
    //     auto traj = valid_trajs[i];
    //     traj.obstacle_score =  traj.obstacle_score / std::max(sum_traj.obstacle_score, 0.01f);
    //     traj.heading_score = traj.heading_score / std::max(sum_traj.heading_score, 0.01f);
    //     traj.speed_score = traj.speed_score / std::max(sum_traj.speed_score, 0.01f);
    //     traj.traj_score = traj.obstacle_score + traj.heading_score + traj.speed_score;
    //     if(traj.traj_score < best_traj_score){
    //         best_traj_score = traj.traj_score;
    //         best_i = i;
    //     }
    // }   
    // 
    cout << "[debug | findBestTrajectory] best_i=" << best_i << " "  << predictedTraj_.size() << endl;
    if(best_i < 0){
        cout << "[debug | findBestTrajectory] can not find best trajectory!" << endl;
        return false;
    }
    best_traj = predictedTraj_[best_i];
    // cout << "[debug | findBestTrajectory] best_i=" << best_i
    //     << ", best_traj_score: " << best_traj_score
    //     << ", best_vel=" << msg_convert::Velocity2D2Str(best_traj.cmd_vel)
    //     << endl;
    publishBestTraj(best_traj);
    return true;    

}


// 计算速度
mc_status DWAController::computeTwist(geometry_msgs::PoseStamped local_goal, geometry_msgs::Twist& cmd_vel){
    mc_status status;
    // 寻找最佳轨迹
    Trajectory best_traj;
    status.success = findBestTrajectory(best_traj, local_goal);
    if(status.success){
        cmd_vel = best_traj.cmd_vel;
    }
    return status;
}